const std = @import("std");
const mem = std.mem;
const fs = std.fs;

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();

    const allocator = gpa.allocator();

    var args = try std.process.argsWithAllocator(allocator);
    defer args.deinit();

    _ = args.skip();

    const root_dir: []const u8 = args.next() orelse return error.MissingTargetDir;

    const addr = try std.net.Address.parseIp("0.0.0.0", 8080);
    var server = try std.net.Address.listen(addr, .{ .reuse_address = true });
    defer server.deinit();

    while (true) {
        var buf: [65535]u8 = undefined;
        const conn = try server.accept();

        var client = std.http.Server.init(conn, &buf);

        while (client.state == .ready) {
            var request = client.receiveHead() catch |err| switch (err) {
                std.http.Server.ReceiveHeadError.HttpConnectionClosing => {
                    std.debug.print("{any}", .{err});
                    return err;
                },
                else => {
                    std.debug.print("{any}", .{err});
                    return err;
                },
            };
            _ = try request.reader();

            const request_info = try parseHttpRequest(&buf);
            const file_path = try buildFilePath(root_dir, request_info.path);
            std.debug.print("{s}\n", .{file_path});

            var is_file = true;
            _ = std.fs.cwd().statFile(file_path) catch |err| {
                if (err == error.FileNotFound) {
                    try request.respond("404", .{});
                    continue;
                }
                is_file = false;
            };

            if (is_file) {
                var file = try fs.openFileAbsolute(file_path, .{ .mode = fs.File.OpenMode.read_only });
                defer file.close();
                var buffer_reader = std.io.bufferedReader(file.reader());
                var reader = buffer_reader.reader();

                var send_buffer: [65536]u8 = undefined;
                const headers = [_]std.http.Header{.{ .name = "Content-Type", .value = "application/octet-stream" }};
                var response = request.respondStreaming(.{
                    .send_buffer = &send_buffer,
                    .respond_options = .{ .extra_headers = &headers },
                });

                var line_buffer: [2048]u8 = undefined;
                while (true) {
                    const size = try reader.read(&line_buffer);
                    if (0 == size) {
                        break;
                    } else {
                        _ = try response.write(line_buffer[0..size]);
                    }
                }
                try response.flush();
                try response.end();
            } else {
                var dir = fs.openDirAbsolute(file_path, .{ .iterate = true }) catch {
                    try request.respond("505", .{});
                    continue;
                };
                defer dir.close();

                var send_buffer: [65536]u8 = undefined;
                var response = request.respondStreaming(.{ .send_buffer = &send_buffer });

                _ = try response.write("<html>\n<body>\n");
                var iterate = dir.iterate();
                while (try iterate.next()) |item| {
                    const download_path = try std.fmt.allocPrint(allocator, "<a href=\"{s}/{s}\">{s}</a><br>\n", .{ request_info.path[1..], item.name, item.name });
                    defer allocator.free(download_path);

                    _ = try response.write(download_path);
                }
                _ = try response.write("</body>\n</html>\n");

                try response.flush();

                try response.end();
            }
        }
    }
}

const RequestInfo = struct { method: []const u8, path: []const u8 };

fn parseHttpRequest(request: []u8) !RequestInfo {
    var lines = std.mem.split(u8, request, "\r\n");
    const request_line = lines.next() orelse return error.InvalidRequest;
    var parts = std.mem.split(u8, request_line, " ");
    const method = parts.next() orelse return error.InvalidRequest;
    const path = parts.next() orelse return error.InvalidRequest;

    return RequestInfo{ .method = method, .path = path };
}

fn buildFilePath(root_dir: []const u8, request_path: []const u8) ![]const u8 {
    // 防止路径遍历攻击
    if (mem.indexOf(u8, request_path, "..") != null) {
        return error.InvalidPath;
    }

    // 组合根目录和请求路径
    return try fs.path.join(std.heap.page_allocator, &[_][]const u8{ root_dir, request_path[1..] });
}

fn getMimeType(path: []const u8) []const u8 {
    if (mem.endsWith(u8, path, ".html")) return "text/html";
    if (mem.endsWith(u8, path, ".css")) return "text/css";
    if (mem.endsWith(u8, path, ".js")) return "application/javascript";
    if (mem.endsWith(u8, path, ".json")) return "application/json";
    if (mem.endsWith(u8, path, ".png")) return "image/png";
    if (mem.endsWith(u8, path, ".jpg") or mem.endsWith(u8, path, ".jpeg")) return "image/jpeg";
    if (mem.endsWith(u8, path, ".gif")) return "image/gif";
    if (mem.endsWith(u8, path, ".svg")) return "image/svg+xml";
    return "application/octet-stream";
}
